﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace NaiveBayesClassifier
{
    public class CNaiveBayes : IClassifier
    {
        private string name;

        protected CTrainingSet trainingSet;
        protected Dictionary<IConcept, Double> conceptPriors;
        protected Dictionary<IConcept, Dictionary<IAttribute, CAttributeValue>> p;
        protected List<String> attributeList;

        public bool verbose = false;

        public CNaiveBayes(string name, CTrainingSet trainingSet)
        {
            this.name = name;
            this.trainingSet = trainingSet;

            conceptPriors = new Dictionary<IConcept, double>(trainingSet.GetNumberOfConcepts());
            verbose = false;
        }

        public bool Train()
        {
            bool hasTrained = false;

            if (attributeList == null || attributeList.Count == 0)
            {
                Console.WriteLine("Can't train the classifier without specifying the attributes for training!");
                Console.WriteLine("Use the method --> trainOnAttribute(Attribute a)");
            }
            else
            {
                CalculateConceptPriors();
                CalculateConditionalProbabilities();
                hasTrained = true;
            }

            return hasTrained;
        }

        private void CalculateConditionalProbabilities()
        {
            p = new Dictionary<IConcept, Dictionary<IAttribute, CAttributeValue>>();
            foreach (IInstance instance in trainingSet.GetInstanceSet())
            {
                foreach (IAttribute attribute in instance.GetAttributes())
                {
                    IConcept concept = instance.GetConcept();
                    if (attribute != null && attributeList.Contains(attribute.GetName()))
                    {
                        if (!p.ContainsKey(concept))
                            p[concept] = new Dictionary<IAttribute, CAttributeValue>();
                        if (p[concept].ContainsKey(attribute))
                        {
                            p[concept][attribute].IncreaseCount();
                        }
                        else
                        {
                            p[concept].Add(attribute, new CAttributeValue(attribute.GetValue()));
                        }
                    }
                }
            }

            //foreach (IConcept concept in p.Keys)
            //{
            //    Console.WriteLine(concept.GetName());
            //    Dictionary<IAttribute, CAttributeValue> map = p[concept];
            //    foreach (IAttribute a in map.Keys)
            //    {
            //        Console.WriteLine(a.GetName() + " " + map[a].ToString() + " " + map[a].GetCount());
            //    }
            //}
        }

        private void CalculateConceptPriors()
        {
            foreach (IConcept concept in trainingSet.GetConceptSet())
                conceptPriors[concept] = 0.0;
            foreach (IInstance instance in trainingSet.GetInstanceSet())
                conceptPriors[instance.GetConcept()]++;
        }

        public IConcept Classify(IInstance instance)
        {
            IConcept bestConcept = null;
            double bestProbability = 0.0;

            if (trainingSet == null || trainingSet.GetNumberOfConcepts() == 0)
                throw new Exception("You have to train classifier first.");

            if (verbose)
                Console.WriteLine("\n*** Classifying instance: {0}\n", instance.ToString());

            foreach (IConcept concept in trainingSet.GetConceptSet())
            {
                double p = GetProbability(concept, instance);
                if (verbose)
                    Console.WriteLine("P({0}|{1}) = {2}\n", concept.GetName(), instance.ToString(), p);
                if (p >= bestProbability)
                {
                    bestProbability = p;
                    bestConcept = concept;
                }
            }

            return bestConcept;
        }

        private double GetProbability(IConcept concept, IInstance instance)
        {
            double conceptProbability = 0.0;
            if (trainingSet.GetConceptSet().Contains(concept))
                conceptProbability = (GetProbability(instance, concept) * GetProbability(concept))
                    / GetProbability(instance);
            else
                conceptProbability = 1.0 / (trainingSet.GetNumberOfConcepts() + 1);
            return conceptProbability;
        }

        private double GetProbability(IInstance instance)
        {
            double conceptProbility = 0.0;
            foreach (IConcept concept in trainingSet.GetConceptSet())
                conceptProbility += GetProbability(instance, concept) * GetProbability(concept);
            if (conceptProbility == 0.0)
                return 1.0 / trainingSet.GetSize();
            else
                return conceptProbility;
        }

        private double GetProbability(IConcept concept)
        {
            double instanceCount;
            if (conceptPriors.ContainsKey(concept))
                instanceCount = conceptPriors[concept];
            else
                instanceCount = 0.0;
            return instanceCount / trainingSet.GetSize();
        }

        private double GetProbability(IInstance instance, IConcept concept)
        {
            double conceptProbability = 1.0;
            foreach (IAttribute attribute in instance.GetAttributes())
            {
                if (attribute != null && attributeList.Contains(attribute.GetName()))
                {
                    if (!p[concept].ContainsKey(attribute))
                        conceptProbability *= (1.0 / (trainingSet.GetSize() + 1));
                    else
                        conceptProbability *= (p[concept][attribute].GetCount() / conceptPriors[concept]);
                }
            }

            if (conceptProbability >= 1.0)
                return 1.0 / trainingSet.GetNumberOfConcepts();
            else
                return conceptProbability;
        }

        public void TrainOnAttribute(string attributeName)
        {
            if (attributeList == null)
                attributeList = new List<string>();
            attributeList.Add(attributeName);
        }

        #region Getter and Setter
        public string GetName()
        {
            throw new NotImplementedException();
        }

        public CTrainingSet GetTrainingSet()
        {
            return trainingSet;
        }
        #endregion
    }
}
